# -*- coding: utf-8 -*-  
'''
字典

@author: luoyi
Created on 2021年3月18日
'''
import pickle

import utils.conf as conf


#    词 -> 索引    初始文案: <PAD>填充字符 <UNK>未知字符 <GO>开始字符 <EOS>结束字符
word_to_index = {'<PAD>':0, '<UNK>':1, '<GO>':2, '<EOS>':3}
#    索引 -> 词
index_to_word = {0:'<PAD>', 1:'<UNK>', 2:'<GO>', 3:'<EOS>'}


#    取词的index（如果word_to_index里有这个词，则直接返回索引，否则扩充字典返回新的索引）
def word2idx(word):
    idx = word_to_index.get(word)
    #    如果有就直接返回
    if (idx is not None): return idx
    
    #    否则扩充字典
    dict_size = len(word_to_index)
    word_to_index[word] = dict_size
    index_to_word[dict_size] = word
    return dict_size


#    取词的index（如果word_to_index里有这个词，则直接返回索引，否则扩充字典返回新的索引）
def word2idx_slist(word_list):
    idx_list = [word2idx(word) for word in word_list]
    return idx_list


#    字典长度
def dict_size():
    return len(word_to_index)


#    加载外部字典
def load_dict_from_file(word2dict_path=conf.TEXT.get_dictionaries_path() + '/word2dict.pkl',
                        dict2word_path=conf.TEXT.get_dictionaries_path() + '/dict2word.pkl'):
    global word_to_index
    global index_to_word
    with open(word2dict_path, mode='rb') as file_word2dict:
        word_to_index = pickle.load(file_word2dict)
    with open(dict2word_path, mode='rb') as file_dict2word:
        index_to_word = pickle.load(file_dict2word)
    pass
    